from func_block import function_block_repo_path
from random_tools.utils import FileUtils
from random_tools.similarity_scores import BM25Search, UniXcoderSearch

from collections import defaultdict
from tqdm import tqdm
import os


def similar_function_blocks_path(base_dir: str, file_name: str, *, ext: str = "pkl"):
    return os.path.join(base_dir, f"{file_name}_sim_func_blocks.{ext}")


def similar_code_blocks_path(base_dir: str, file_name: str, *, ext: str = "pkl"):
    return os.path.join(base_dir, f"{file_name}_sim_code_blocks.{ext}")


class SimilarFunctionQuery:
    def __init__(self, *, repo_base: str,
                 repo_block_base: str,
                 query_block_dataset_path: str,
                 sim_func_path: str,
                 save_base: str,
                 top_k_func: int = 2,
                 top_k_block: int = 10,
                 retriever_type: str = "bm25",
                 model_path: str | None = None,
                 device: str = "cpu"):
        self.repo_base = repo_base
        self.repo_block_base = repo_block_base
        self.query_block_dataset = query_block_dataset_path
        self.sim_func_path = sim_func_path
        self.save_base = save_base

        self.top_k_func = top_k_func
        self.top_k_block = top_k_block

        self.retriever_type = retriever_type
        if retriever_type == "unixcoder":
            self.retriever = UniXcoderSearch(model_path, device=device)

    def search_for_dataset(self, *, save_name: str, save_pickle: bool = True, save_json: bool = False):
        sim_func_path = FileUtils.load_pickle(self.sim_func_path)
        query_block_dataset = FileUtils.load_pickle(self.query_block_dataset)

        sim_func_blocks_all: dict[str, list[dict]] = defaultdict(list)
        sim_code_blocks_all: dict[str, list[dict]] = defaultdict(list)

        for data in tqdm(query_block_dataset):
            metadata = data["metadata"]
            task_id = metadata["task_id"]
            line_no = metadata["line_no"]
            similar_functions = sim_func_path.get(task_id, [])

            fpath_tuple = tuple(metadata["fpath_tuple"])
            repo_name = fpath_tuple[0]

            repo_blocks = FileUtils.load_pickle(function_block_repo_path(self.repo_block_base, repo_name))
            filtered_candidates: list[tuple[tuple[str, ...], str, str | None]] = []

            for candidate_func in similar_functions:
                cand_fpath = candidate_func["fpath_tuple"]
                cand_func_name = candidate_func["function_name"]
                cand_class_method = candidate_func["class_method"]

                filtered_candidates.append(
                    (cand_fpath, cand_func_name, cand_class_method)
                )

                cand_blocks = [
                    block for block in repo_blocks
                    if (
                            block["metadata"]["fpath_tuple"] == cand_fpath
                            and block["metadata"]["function_name"] == cand_func_name
                            and block["metadata"]["class_method"] == cand_class_method
                    )
                ]

                if len(cand_blocks) == 0:
                    continue

                if self.retriever_type == "unixcoder":
                    results = self.retriever.query(data["context"], cand_blocks, top_k=self.top_k_func)
                    sim_func_blocks_all[task_id] += results
                else:
                    bm25_search = BM25Search(cand_blocks)
                    results = bm25_search.query(data["context"], top_k=self.top_k_func)
                    sim_func_blocks_all[task_id] += results

            filtered_blocks = [
                block for block in repo_blocks
                if (block["metadata"]["fpath_tuple"],
                    block["metadata"]["function_name"],
                    block["metadata"]["class_method"])
                   not in filtered_candidates
            ]

            filtered_blocks = [
                block for block in filtered_blocks
                if not (
                        block["metadata"]["fpath_tuple"] == tuple(metadata["fpath_tuple"])
                        and block["metadata"]["end_line_no"] >= metadata["context_start_lineno"]
                )
            ]

            filtered_blocks_text = [
                block["context"] for block in filtered_blocks
            ]

            if self.retriever_type == "unixcoder":
                results = self.retriever.query(data["context"], filtered_blocks_text, top_k=self.top_k_block)
            else:
                bm25_search = BM25Search(filtered_blocks)
                results = bm25_search.query(data["context"], top_k=self.top_k_block)
            sim_code_blocks_all[task_id] += results

        if save_pickle:
            FileUtils.save_pickle(sim_func_blocks_all,
                                  similar_function_blocks_path(self.save_base, save_name, ext="pkl"))
            FileUtils.save_pickle(sim_code_blocks_all,
                                  similar_code_blocks_path(self.save_base, save_name, ext="pkl"))

        if save_json:
            FileUtils.save_json(sim_func_blocks_all,
                                similar_function_blocks_path(self.save_base, save_name, ext="json"))
            FileUtils.save_json(sim_code_blocks_all,
                                similar_code_blocks_path(self.save_base, save_name, ext="json"))

